package bloomfilter

import (
	"context"
	"github.com/redis/go-redis/v9"
	"hash/fnv"
	"math"
)

type RedisBloomFilter struct {
	rdb        *redis.Client // Redis 客户端
	key        string        // Redis 中存储位图的键名
	numHashes  int           // 哈希函数数量
	bitmapSize uint32        // 位图大小（总位数）
}

// NewRedisBloomFilter 创建布隆过滤器实例
func NewRedisBloomFilter(rdb *redis.Client, key string, capacity int, errorRate float64) *RedisBloomFilter {
	// 计算最优位图大小和哈希函数数量
	bitmapSize := calculateBitmapSize(capacity, errorRate)
	numHashes := calculateNumHashes(bitmapSize, capacity)

	return &RedisBloomFilter{
		rdb:        rdb,
		key:        key,
		numHashes:  numHashes,
		bitmapSize: uint32(bitmapSize),
	}
}

// calculateBitmapSize 计算位图大小（公式：m = - (n * ln(p)) / (ln(2)^2)）
func calculateBitmapSize(n int, p float64) uint32 {
	size := -float64(n) * math.Log(p) / math.Pow(math.Log(2), 2)
	return uint32(math.Ceil(size))
}

// calculateNumHashes 计算哈希函数数量（公式：k = (m / n) * ln(2)）
func calculateNumHashes(m uint32, n int) int {
	k := float64(m) / float64(n) * math.Log(2)
	return int(math.Ceil(k))
}

// getPositions 计算元素在多个哈希函数下的位位置
func (bf *RedisBloomFilter) getPositions(data string) []uint32 {
	positions := make([]uint32, bf.numHashes)
	h := fnv.New32a() // 使用FNV哈希算法

	for i := 0; i < bf.numHashes; i++ {
		h.Reset()
		h.Write([]byte(data))
		h.Write([]byte{byte(i)}) // 添加种子值生成不同哈希
		hash := h.Sum32()
		positions[i] = hash % bf.bitmapSize
	}
	return positions
}

// Add 添加元素到位图
func (bf *RedisBloomFilter) Add(ctx context.Context, data string) error {
	positions := bf.getPositions(data)
	pipe := bf.rdb.Pipeline()
	for _, pos := range positions {
		pipe.SetBit(ctx, bf.key, int64(pos), 1) // SETBIT 设置位
	}
	_, err := pipe.Exec(ctx)
	return err
}

// Exists 检查元素是否存在
func (bf *RedisBloomFilter) Contains(ctx context.Context, data string) (bool, error) {
	positions := bf.getPositions(data)
	pipe := bf.rdb.Pipeline()
	for _, pos := range positions {
		pipe.GetBit(ctx, bf.key, int64(pos)) // GETBIT 获取位值
	}
	results, err := pipe.Exec(ctx)
	if err != nil {
		return false, err
	}

	// 检查所有位是否均为1
	for _, res := range results {
		if res.(*redis.IntCmd).Val() == 0 {
			return false, nil // 任意一位为0则一定不存在
		}
	}
	return true, nil // 所有位为1时可能存在（有误判概率）
}
